import asyncio
import json
import time
from typing import Dict

from common import *

VERSION = 'v1.0.0'

log = get_logger('agent')
log.info("agent version: %s", VERSION)

class Session:
    def __init__(self, agent, session_id, target_addr):
        self.agent = agent
        self.session_id = session_id
        self.target_addr = target_addr
        self.target_connected = False
        self.client_closed = False
        self.target_reader = None
        self.target_writer = None
        self.tunnel_queue = asyncio.Queue()
        self.status = 'connecting'  # 会话状态：connecting, connected, closed

    def close(self):
        if self.status == 'closed':
            return
        self.status = 'closed'
        # 1. close target side
        if self.target_connected and self.target_writer:
            log.info(f"[{self.session_id}] closing target")
            self.target_connected = False
            try:
                self.target_writer.close()
            except Exception as e:
                log.error(f"[{self.session_id}] close target error: {e}")
        # 2. close client side
        self.tunnel_queue.put_nowait(None)
        self.agent.remove_session(self.session_id)

    async def connect_target(self, target_addr):
        """连接到target并启动转发任务"""
        if not target_addr:
            log.error(f"会话{self.session_id}未指定target地址")
            return
        try:
            target_host, target_port = target_addr.split(':')
            self.target_reader, self.target_writer = await asyncio.open_connection(target_host, int(target_port))
            self.target_connected = True
            log.info(f"[{self.session_id}] conntected to target: {target_addr}")
            # 发送连接确认
            self.agent.writer.write(encode_tlv(MSG_TYPE_SESSION_CONNECTED, self.session_id.encode()))
            await self.agent.writer.drain()
            log.info(f"[{self.session_id}] sent connected message to tunnel server")
            # 启动双向转发任务
            asyncio.create_task(self.agent_to_target())
            asyncio.create_task(self.target_to_agent())
        except Exception as e:
            log.error(f"会话{self.session_id}无法连接到target {target_addr}: {e}")
            self.agent.writer.write(encode_tlv(MSG_TYPE_SESSION_CLOSE, self.session_id.encode()))
            await self.agent.writer.drain()
            log.info(f"[{self.session_id}] 已发送会话关闭消息")

    async def agent_to_target(self):
        """从隧道队列读取数据并转发到target"""
        try:
            while self.target_connected:
                data = await self.tunnel_queue.get()
                if not data:
                    log.info(f"[{self.session_id}][c->s] tunnel read queue closed")
                    break
                self.target_writer.write(data)
                await self.target_writer.drain()
                self.tunnel_queue.task_done()
                # print(f"[{self.session_id}] tunnel -> target: {len(data)} bytes")
        except Exception as e:
            log.error(f"[{self.session_id}][c->s] error: {e}")
        finally:
            self.close()
        log.info(f"[{self.session_id}][c->s] loop exited.")

    async def target_to_agent(self):
        """从target读取响应数据并转发到server"""
        try:
            while self.target_connected:
                data = await self.target_reader.read(4096)
                if not data:
                    log.info(f'[{self.session_id}][s->c] target closed')
                    break
                # 封装为TLV格式发送回server
                tun_data = self.session_id.encode() + data
                tlv_data = encode_tlv(MSG_TYPE_SESSION_DATA, tun_data)
                self.agent.writer.write(tlv_data)
                await self.agent.writer.drain()
                # print(f"[{self.session_id}] tunnel <- target: {len(data)} bytes")
        except Exception as e:
            log.error(f"[{self.session_id}][s->c] error: {e}")
        finally:
            self.close()
        log.info(f"[{self.session_id}][s->c] loop exited.")

class Agent:
    def __init__(self, agent_id, server_host='localhost', server_port=8889):
        self.agent_id = agent_id
        self.server_host = server_host
        self.server_port = server_port
        self.reader = None
        self.writer = None
        self.config = {}
        self.running = False
        self.dispatcher_task = None
        self.heartbeat_task = None
        self.retry_interval = 10  # 初始重连间隔(秒)
        self.sessions: Dict[str, Session] = {}

    def remove_session(self, session_id):
        if session_id in self.sessions:
            del self.sessions[session_id]
            log.info(f"[{session_id}] remove session")

    async def connect(self):
        while True:
            try:    
                log.info(f'Connecting to server {self.server_host}:{self.server_port}')
                self.reader, self.writer = await asyncio.open_connection(self.server_host, self.server_port)
                log.info(f'Connected. sending register request')
                # 发送agent注册信息
                register_data = json.dumps({'agent_id': self.agent_id, 'version': VERSION}).encode()
                self.writer.write(encode_tlv(MSG_TYPE_REGISTER, register_data))
                await self.writer.drain()
                # 接收服务器响应
                type_, value = await decode_tlv(self.reader)
                if type_ != MSG_TYPE_REGISTER_RESP: 
                    raise TLVError(f"Expected register response (type {MSG_TYPE_REGISTER_RESP}), got {type_}")
                response = json.loads(value.decode())
                log.info(f'register response: {response}')
                return response
            except Exception as e:
                log.error(f'[connect] error: {e}. retry after {self.retry_interval} seconds')
                await asyncio.sleep(self.retry_interval)

    async def run(self):
        self.running = True
        while self.running:
            try:
                response = await self.connect()

                if 'config' in response:
                    self.config = response['config']
                    self.apply_config()

                # 启动消息分发和心跳任务
                self.dispatcher_task = asyncio.create_task(self.message_dispatcher())
                self.heartbeat_task = asyncio.create_task(self.task_heartbeat())

                # 等待任务完成
                await asyncio.gather(self.dispatcher_task, self.heartbeat_task)
            except asyncio.CancelledError:
                print(f'[cancle] user cancelled program')
                self.stop()
            except Exception as e:
                print(f'[cancle] unexpected error: {e}')
            print('[cancle] cancel tasks safely...')
            try:
                if self.dispatcher_task and not self.dispatcher_task.done():
                    self.dispatcher_task.cancel()
                if self.heartbeat_task and not self.heartbeat_task.done():
                    self.heartbeat_task.cancel()
                # 等待任务取消完成
                if self.dispatcher_task:
                    await asyncio.gather(self.dispatcher_task, self.heartbeat_task, return_exceptions=True)
            except Exception as e:
                print(f'[cancle] Task cancellation error: {e}')
            # 主动关闭连接
            print('[cancle] closing tunnel connection...')
            if self.writer:
                try:
                    self.writer.close()
                    await self.writer.wait_closed()
                except Exception as e:
                    print(f'[cancle] Connection close error: {e}')
            if self.running:
                print(f'[reconnecting] retry after {self.retry_interval} seconds')
                await asyncio.sleep(self.retry_interval)
        print('[exit] agent stopped')

    def stop(self):
        """停止agent并退出重连循环"""
        self.running = False

    async def message_dispatcher(self):
        """分发接收到的消息"""
        log.info('[dispatcher] start loop...')
        while self.running:
            try:
                type_, value = await decode_tlv(self.reader)
                if type_ == MSG_TYPE_CONFIG_UPDATE:
                    await self.handle_config_update(value)
                    continue
                elif type_ == MSG_TYPE_HEARTBEAT:
                    # heartbeat_info = json.loads(value.decode())
                    # print(f"收到来自server的心跳: {heartbeat_info}")
                    continue
                # following msg need session_id
                if len(value) < 8:
                    log.warning(f"warn: missing session_id: {value}")
                    continue
                session_id = value[:8].decode()
                payload = value[8:] if len(value) > 8 else b''
                if type_ == MSG_TYPE_SESSION_NEW:
                    await self.handle_session_new(session_id, payload)
                elif type_ == MSG_TYPE_SESSION_DATA:
                    await self.handle_tunnel_data(session_id, payload)
                elif type_ == MSG_TYPE_SESSION_CLOSE:
                    log.info(f"[{session_id}] recv close msg from tun-server.")
                    if session_id in self.sessions:
                        session = self.sessions[session_id]
                        del self.sessions[session_id]
                        session.client_closed = True
                        session.close()
                        log.info(f"[{session_id}] session removed")
                    else:
                        log.warning(f"[{session_id}] session not found.")
                else:
                    log.warning(f"未知消息类型: {type_}")
            except TLVError as e:
                log.warning(f"[dispatcher] read tlv error: {e}.")
                break
            except Exception as e:
                log.error(f"[dispatcher] error: {e}")
                raise
        log.info('[dispatcher] loop exited.')

    async def task_heartbeat(self):
        """定期发送保活消息到server"""
        log.info('[heartbeat] start loop...')

        while self.running:
            try:
                heartbeat_data = json.dumps({'total_session': len(self.sessions), 'timestamp': time.time()}).encode()
                self.writer.write(encode_tlv(MSG_TYPE_HEARTBEAT, heartbeat_data))
                await self.writer.drain()
                await asyncio.sleep(10)
            except Exception as e:
                log.error(f"[heartbeat] error: {e}.")
                log.info('[heartbeat] cancel dispatcher task...')
                if self.dispatcher_task and not self.dispatcher_task.done():
                    self.dispatcher_task.cancel()
                    await asyncio.gather(self.dispatcher_task, return_exceptions=True)
                # todo: close all session connection
                for sid, session in self.sessions.items():
                    log.info(f'[{sid}] closing session connection. target: {session.target_addr}')
                    if session.target_writer:
                        session.target_writer.close()
                        # await session.target_writer.wait_closed()
                log.info('[heartbeat] reconnecting to tunnel server...')
                await self.connect()
                log.info('[heartbeat] reconnected. re-create dispatcher task...')
                self.dispatcher_task = asyncio.create_task(self.message_dispatcher())
        log.info('[heartbeat] task stopped')

    async def handle_config_update(self, value):
        """处理配置更新消息"""
        update = json.loads(value.decode())
        print(f'Received config update: {update}')
        if 'config' in update:
            self.config = update['config']
            self.apply_config()

    async def handle_session_new(self, session_id, payload):
        """处理新会话请求"""
        try:
            session_data = json.loads(payload.decode())
            target_addr = session_data.get('target')
            if not target_addr:
                log.warning(f"[{session_id}] new消息缺少target")
                return
            log.info(f"[{session_id}] new conn. target: {target_addr}")
            session = Session(self, session_id, target_addr)
            # 连接到target
            await session.connect_target(target_addr)
            if session.target_connected:
                log.info(f"[{session_id}] add session")
                self.sessions[session_id] = session
        except Exception as e:
            log.error(f"[{session_id}] handle_new error: {e}")

    async def handle_tunnel_data(self, session_id, payload):
        """处理隧道数据并转发到target"""
        session = self.sessions.get(session_id)
        if not session:
            log.error(f"[{session_id}] does not exist.")
            return
        # 将数据放入隧道队列
        await session.tunnel_queue.put(payload)

    def apply_config(self):
        log.info(f'Applying new configuration: {self.config}')
        # 这里添加配置应用逻辑

async def main():
    import sys
    if len(sys.argv) != 2:
        print('Usage: python agent.py <ip>:<port>/<agent_id>')
        sys.exit(1)

    agent_id = sys.argv[1].split('/')[-1]
    server_host, server_port = sys.argv[1].split('/')[0].split(':')
    agent = Agent(agent_id, server_host, int(server_port))
    await agent.run()

if __name__ == '__main__':
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        log.info('Agent shutdown')